import os
import argparse

import cv2
import numpy as np
from tqdm import tqdm
import sys

sys.path.append('../')  # 确保这个路径是正确的
from cyw_devkit.core.dataset import CywDataset, transform_bboxes
from cyw_devkit.core.visualizer import Visualizer, tranform_pt_pad
from tictoc import TicToc
import multiprocessing
import pickle
from cyw_devkit.core.io import read_points

def main(args):
    # 统计耗时
    cost = TicToc("数据可视化")
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    files = []
    for dir in frames:
        if dir[:2] == '__':  # '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                files.append(dir)

    process_size = len(files)
    # for idx in range(process_size):
    #     main_worker(files[0], args)

    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(dataset_path, args):
    dataset = pickle.load(open(os.path.join(dataset_path,'data_base/infos_trainval.pkl'), 'rb'))
    for idx, (stamp, data) in enumerate(dataset.items()):
        if idx < 17:
            continue
        lidar = read_points(os.path.join(args.data_path,data['lidar']['path']), 4)
        for sweep in data['sweeps']:
            lidar_pre = read_points(os.path.join(args.data_path,sweep['lidar']['path']), 4)
            pre2cur = sweep['pre2cur']
            # pre2cur[2,3] = 0.0
            lidar2cur = tranform_pt_pad(lidar_pre, tfs=[pre2cur])
            lidar = np.concatenate((lidar, lidar2cur), axis=0)
        boxes = np.concatenate((data['label']['location'],data['label']['dimensions'],data['label']['rotation_y']),axis=1)
        vis = Visualizer(lidar, center_mode=False)
        vis.add_bboxes(bbox3d=boxes)
        vis.show()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default="/home/adt/bags/work_space/datasets",
                        help='your data root')
    args = parser.parse_args()

    main(args)
